
# 这行代码将不在 PyFlow 目录内安装的扩展包添加进来
__path__ = __import__("pkgutil").extend_path(__path__, __name__)

# 导入 importlib 模块，用于动态导入模块
import importlib
# 导入 pkgutil 模块，用于查找模块和包
import pkgutil
# 从 collections.abc 导入抽象基类，用于类型检查
import collections.abc
# 从 copy 模块导入 copy 函数，用于浅拷贝对象
from copy import copy
# 导入 os 模块，用于与操作系统进行交互，处理文件路径等
import os
# 导入 json 模块，用于处理 JSON 数据
import json

# 从 PyFlow.Packages 包导入所有内容
from src.Packages import *


# 定义模块对外暴露的接口列表
__all__ = [
    "INITIALIZE",
    "GET_PACKAGE_CHECKED",
    "GET_PACKAGE_PATH",
    "GET_PACKAGES",
    "CreateRawPin",
    "getPinDefaultValueByType",
    "findPinClassByType",
    "getRawNodeInstance",
    "getAllPinClasses",
    "getHashableDataTypes",
]


# 存储所有已加载的包的字典
__PACKAGES = {}
# 存储每个包对应的路径的字典
__PACKAGE_PATHS = {}
# 存储可哈希数据类型的列表
__HASHABLE_TYPES = []


def GET_PACKAGES():
    """
    返回所有已加载的包的字典。

    :return: 包含所有已加载包的字典
    """
    return __PACKAGES


def GET_PACKAGE_PATH(packageName):
    """
    根据包名返回对应的包路径。

    :param packageName: 包的名称
    :return: 包的路径，如果包名不存在则返回 None
    """
    if packageName in __PACKAGE_PATHS:
        return __PACKAGE_PATHS[packageName]


def GET_PACKAGE_CHECKED(package_name):
    """
    确保指定的包已加载，并返回该包。

    :param package_name: 包的名称
    :return: 已加载的包对象
    :raises AssertionError: 如果指定的包未加载
    """
    assert package_name in __PACKAGES
    return __PACKAGES[package_name]


def getAllPinClasses():
    """
    获取所有包中定义的引脚类列表。

    :return: 包含所有引脚类的列表
    """
    result = []
    for package in list(__PACKAGES.values()):
        result += list(package.GetPinClasses().values())
    return result


def findPinClassByType(dataType):
    """
    根据数据类型查找对应的引脚类。

    :param dataType: 数据类型
    :return: 对应的引脚类，如果未找到则返回 None
    """
    for package_name, package in GET_PACKAGES().items():
        pins = package.GetPinClasses()
        if dataType in pins:
            return pins[dataType]
    return None


def getPinDefaultValueByType(dataType):
    """
    根据数据类型获取引脚的默认值。

    :param dataType: 数据类型
    :return: 引脚的默认值，如果未找到对应的引脚类则返回 None
    """
    pin = findPinClassByType(dataType)
    if pin:
        return pin.pinDataTypeHint()[1]
    return None


def getHashableDataTypes():
    """
    获取所有可哈希的数据类型列表。

    :return: 包含可哈希数据类型名称的列表的浅拷贝
    """
    if len(__HASHABLE_TYPES) == 0:
        for pin in getAllPinClasses():
            t = pin.internalDataStructure()
            if t is not type(None) and t is not None:
                if isinstance(pin.pinDataTypeHint()[1], collections.abc.Hashable):
                    __HASHABLE_TYPES.append(pin.__name__)
    return copy(__HASHABLE_TYPES)


def getPinFromData(data):
    """
    根据数据查找对应的引脚类。

    :param data: 数据
    :return: 对应的引脚类，如果未找到则返回 None
    """
    for pin in [pin for pin in getAllPinClasses() if pin.IsValuePin()]:
        pType = pin.internalDataStructure()
        if data == pType:
            return pin


def CreateRawPin(name, owningNode, dataType, direction, **kwds):
    """
    创建一个原始的引脚实例。

    :param name: 引脚的名称
    :param owningNode: 拥有该引脚的节点
    :param dataType: 引脚的数据类型
    :param direction: 引脚的方向
    :param kwds: 其他关键字参数
    :return: 引脚实例，如果未找到对应的引脚类则返回 None
    """
    pinClass = findPinClassByType(dataType)
    if pinClass is None:
        return None
    inst = pinClass(name, owningNode, direction, **kwds)
    return inst


def getRawNodeInstance(nodeClassName, packageName=None, libName=None, **kwargs):
    """
    根据节点类名、包名和库名获取原始节点实例。

    :param nodeClassName: 节点类的名称
    :param packageName: 包的名称，默认为 None
    :param libName: 库的名称，默认为 None
    :param kwargs: 其他关键字参数
    :return: 节点实例，如果未找到则尝试不同方式查找并返回对应的节点实例
    """
    from src.Core.NodeBase import NodeBase

    package = GET_PACKAGE_CHECKED(packageName)
    # 先尝试查找函数节点
    if libName is not None:
        for key, lib in package.GetFunctionLibraries().items():
            foos = lib.getFunctions()
            if libName == key and nodeClassName in foos:
                return NodeBase.initializeFromFunction(foos[nodeClassName])

    # 尝试查找节点类
    nodes = package.GetNodeClasses()
    if nodeClassName in nodes:
        return nodes[nodeClassName](nodeClassName, **kwargs)

    # 尝试查找导出的 Python 节点
    packagePath = GET_PACKAGE_PATH(packageName)
    pyNodesPath = os.path.join(packagePath, "PyNodes")
    if os.path.exists(pyNodesPath):
        for path, dirs, files in os.walk(pyNodesPath):
            for pyNodeFileName in files:
                pyNodeName, _ = os.path.splitext(pyNodeFileName)
                if nodeClassName == pyNodeName:
                    pythonNode = getRawNodeInstance("pythonNode", "PyFlowBase")
                    pyNodeFullPath = os.path.join(path, pyNodeFileName)
                    with open(pyNodeFullPath, "r") as f:
                        pythonNode._nodeData = f.read()
                    return pythonNode

    # 尝试查找导出的复合节点
    compoundNodesPath = os.path.join(packagePath, "Compounds")
    if os.path.exists(compoundNodesPath):
        for path, dirs, files in os.walk(compoundNodesPath):
            for compoundNodeFileName in files:
                compoundNodeName, _ = os.path.splitext(compoundNodeFileName)
                compoundNodeFullPath = os.path.join(path, compoundNodeFileName)
                with open(compoundNodeFullPath, "r") as f:
                    compoundData = json.load(f)
                    if compoundData["name"] == nodeClassName:
                        compoundNode = getRawNodeInstance("compound", "PyFlowBase")
                        compoundNodeFullPath = os.path.join(path, compoundNodeFileName)
                        with open(compoundNodeFullPath, "r") as f:
                            jsonString = f.read()
                            compoundNode._rawGraphJson = json.loads(jsonString)
                        return compoundNode


def INITIALIZE(additionalPackageLocations=None, software=""):
    """
    初始化所有包，加载包并注册相关工厂和工具。

    :param additionalPackageLocations: 额外的包路径列表，默认为 None
    :param software: 软件名称，默认为空字符串
    """
    # 清空已加载的包、包路径和可哈希数据类型列表
    __PACKAGES.clear()
    __PACKAGE_PATHS.clear()
    __HASHABLE_TYPES.clear()
    if additionalPackageLocations is None:
        additionalPackageLocations = []
    from src.UI.Tool import REGISTER_TOOL
    from src.UI.Widgets.InputWidgets import REGISTER_UI_INPUT_WIDGET_PIN_FACTORY
    from src.UI.Canvas.UINodeBase import REGISTER_UI_NODE_FACTORY
    from src.UI.Canvas.UIPinBase import REGISTER_UI_PIN_FACTORY
    from src import ConfigManager
    from qtpy.QtWidgets import QMessageBox

    # 获取 Packages 包的路径
    packagePaths = Packages.__path__

    def ensurePackagePath(inPath):
        """
        确保路径指向正确的包路径。

        :param inPath: 输入的路径
        :return: 正确的包路径
        """
        for subFolder in os.listdir(inPath):
            subFolderPath = os.path.join(inPath, subFolder)
            if os.path.isdir(subFolderPath):
                if "PyFlow" in os.listdir(subFolderPath):
                    subFolderPath = os.path.join(subFolderPath, "PyFlow", "Packages")
                    if os.path.exists(subFolderPath):
                        return subFolderPath
        return inPath

    def recursePackagePaths(inPath):
        """
        递归查找指定路径下的所有包路径。

        :param inPath: 输入的路径
        :return: 包含所有包路径的列表
        """
        paths = []
        for subFolder in os.listdir(inPath):
            subFolderPath = os.path.join(inPath, subFolder)
            if os.path.isdir(subFolderPath):
                if "PyFlow" in os.listdir(subFolderPath):
                    subFolderPath = os.path.join(subFolderPath, "PyFlow", "Packages")
                    if os.path.exists(subFolderPath):
                        paths.append(subFolderPath)
        return paths

    # 检查环境变量中的额外包路径
    if "PYFLOW_PACKAGES_PATHS" in os.environ:
        delim = ";"
        pathsString = os.environ["PYFLOW_PACKAGES_PATHS"]
        # 移除路径字符串右侧的分隔符
        pathsString = pathsString.rstrip(delim)
        for packagesRoot in pathsString.split(delim):
            if os.path.exists(packagesRoot):
                paths = recursePackagePaths(packagesRoot)
                packagePaths.extend(paths)

    for packagePathId in range(len(additionalPackageLocations)):
        packagePath = additionalPackageLocations[packagePathId]
        packagePath = ensurePackagePath(packagePath)
        additionalPackageLocations[packagePathId] = packagePath

    # 将额外的包路径添加到包路径列表中
    packagePaths.extend(additionalPackageLocations)

    # 遍历所有包路径，加载包
    for importer, modname, ispkg in pkgutil.iter_modules(packagePaths):
        try:
            if ispkg:
                print("Loading package: {0}".format(modname))
                mod = importer.find_spec(modname).loader.load_module()
                package = getattr(mod, modname)()
                __PACKAGES[modname] = package
                __PACKAGE_PATHS[modname] = os.path.normpath(mod.__path__[0])
        except Exception as e:
            # 加载包出错时弹出错误消息框
            QMessageBox.critical(
                None, "Fatal error", "Error On Module %s :\n%s" % (modname, str(e))
            )
            continue

    # 存储已注册的内部引脚数据类型
    registeredInternalPinDataTypes = set()

    # 遍历所有已加载的包，进行相关注册操作
    for name, package in __PACKAGES.items():
        packageName = package.__class__.__name__
        for node in package.GetNodeClasses().values():
            node._packageName = packageName

        for pin in package.GetPinClasses().values():
            pin._packageName = packageName
            if pin.IsValuePin():
                internalType = pin.internalDataStructure()
                if internalType in registeredInternalPinDataTypes:
                    raise Exception(
                        "Pin from package {0} with {1} internal data type already been registered".format(
                            packageName, internalType
                        )
                    )
                registeredInternalPinDataTypes.add(internalType)

        uiPinsFactory = package.UIPinsFactory()
        if uiPinsFactory is not None:
            REGISTER_UI_PIN_FACTORY(packageName, uiPinsFactory)

        uiPinInputWidgetsFactory = package.PinsInputWidgetFactory()
        if uiPinInputWidgetsFactory is not None:
            REGISTER_UI_INPUT_WIDGET_PIN_FACTORY(packageName, uiPinInputWidgetsFactory)

        uiNodesFactory = package.UINodesFactory()
        if uiNodesFactory is not None:
            REGISTER_UI_NODE_FACTORY(packageName, uiNodesFactory)

        for toolClass in package.GetToolClasses().values():
            supportedSoftwares = toolClass.supportedSoftwares()
            if "any" not in supportedSoftwares:
                if software not in supportedSoftwares:
                    continue
            REGISTER_TOOL(packageName, toolClass)
    # 获取可哈希的数据类型列表
    getHashableDataTypes()